from tkinter import *
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.figure as figure
from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg
import os
from os import path
from msvcrt import getch
from data import Dataset
from utils import Trade_date
from time import time
from multiprocessing.pool import ThreadPool

class Speculate():
    def __init__(self):
        self.df = pd.DataFrame() # 预测数据
        self.mega_days = 30 # 日线单位，计算大趋势
        self.sub_days = 1 # 日线单位，计算小趋势
        self.data = Dataset()
        
    def plt_one_data(self,data,train_days:int,plt:figure.Figure,output:FigureCanvasTkAgg):
        """生成1条数据
        1、使用同长度等长数列作为x轴，后再用日期替换掉
        2、train_days,从最早一天数据开始，后train_days天的数据会保存下来作为训练数据
        3、如果train_days大于数据长度，则跳过这次数据，保持数据一致性
        """
        output.get_tk_widget().pack_forget()
        plt.clear() # 清空画布
        reg,b = self.get_target(data,len(data)-train_days)
        x,y = range(1,data.shape[0]+1),data['close']
        x0,y0 = x[0],x[0]*reg+b
        x1,y1 = x[-1],x[-1]*reg+b
        plt.plot([x0,x1],[y0,y1])
        xlabel = data['trade_date']
        # 画图
        plt.plot(x,y,linewidth=5,marker='o',ms=20)
        # 标记训练数据结束点
        if train_days > len(x):
            return
        plt.text(x[train_days-1]+0.5,y[train_days-1],'<=',ha='center', va='center', fontsize=40,color='r')
        # 显示价格数据
        for a,b in zip(x,y):
            plt.text(a,b,b,ha='center', va='bottom', fontsize=30)
        # 把x轴标签换成日期
        plt.set_xticks(x,xlabel,rotation=-60)
        output.draw()
        output.get_tk_widget().pack()
        
        
    def get_target(self,data,train_days,threshold=3):
        """获取标签
        1、对所有训练的点回归，根据斜率看大趋势是不是向上的
        2、检查后推测那天后几天是涨了还是跌了
        Args:
            data (_type_): 打标签的数据
            train_days (_type_): 用于训练的天数，剩下后面的用于判断涨跌
            threshold (_type_): 涨幅的阈值
        return
            label 1涨 0跌或平
        """
        y = data[:train_days]['close']
        x = range(1,len(y)+1)
        reg,_ = np.polyfit(x,y,1)
        # if not reg>0:
        #     return 0
        py = data[train_days+1:]['pct_chg']
        chg = 0
        for item in py:
            chg += float(item)
            break
        if chg >= threshold:
            return 1
        return 0
        
    
    def save_noe_data(self,file_name,data:pd.DataFrame,target,preprocess):
        """处理保存一个数据
        Args:
            file_name (_type_): 待处理文件名
            data (pd.DataFrame): 待处理数据
            target (_type_): 标签
            preprocess (_type_): 处理进度
        """
        # 从文件名取出结束日期
        end_date = file_name.split('_')[-1][:-4]
        if not end_date.isdigit():
            return
        # 去掉后面的参考数据,然后按日期排好序
        data = data[data['trade_date']<=int(end_date)]
        data = data.sort_values(by='trade_date',ascending=True).reset_index(drop=True)
        if len(data) == 0:
            return
        # 取出close数据,和最后5天数据
        close_data = data['close']
        end5_date_data = data.iloc[-5:].reset_index(drop=True)
        # 保存数据，每天为一列，最后5天成交量,最后1天星期几 各为一列
        save_data = {}
        columns = []
        for i,sub_data in zip(range(len(close_data)),close_data):
            save_data[str(i+1)] = [sub_data]
            columns.append(str(i+1))
        #save_data['pct_chg'] = [end5_date_data['pct_chg'].iloc[0]]
        for i in range(len(end5_date_data)):
            save_data[f'vol{i}'] = [end5_date_data.iloc[i]['vol']]
            columns.append(f'vol{i}')
        save_data['month'] = end_date[4:6]
        save_data['day'] = end_date[-2:]
        save_data['week'] = [pd.to_datetime(end_date,format='%Y%m%d').dayofweek+1]
        save_data['target'] = [target]
        columns = columns + ['month','day','week','target']
        save_data = pd.DataFrame(save_data,columns=columns)
        last_name ='日线_成交量_星期_data/train_data/{}'.format(str(preprocess))+'.csv'
        save_name ='日线_成交量_星期_data/train_data/{}'.format(str(preprocess+1))+'.csv'
        if not path.exists(last_name):
            save_data.to_csv(save_name,index=False)
        else:
            save_data.to_csv(last_name,mode='a',index=False,header=False)
            os.rename(last_name,save_name)
        
    
    def mkdata_by_plt(self,train_days=20):
        """
        1、检查有没有待处理数据
        2、创建interTk窗口，默认大小1080p
        3、设置plt字体
        4、创建interTk画布和plt画布，plt画布放在inerTk画布里
        5、获取数据的生成器，点一下按钮获取下一个数据
        6、train_days: 每个数据训练天数长度
        """
        # 检查有没有数据处理
        if not path.exists('日线_成交量_星期_data/raw_data'):
            print('没有数据可处理')
        file_list = os.listdir('日线_成交量_星期_data/raw_data')
        if len(file_list) == 0:
            print('没有数据可处理')
        # 按照排序读取进度
        file_list.sort()
        self.preprocess = 0
        if not path.exists('日线_成交量_星期_data/train_data'):
            os.mkdir('日线_成交量_星期_data/train_data')
        else: # 获取之前的进度
            train_file = os.listdir('日线_成交量_星期_data/train_data')
            if not len(train_file) == 0:
                train_file.sort()
                preprocess = train_file[-1].split('.')[0]
                if preprocess.isdigit() and str(len(file_list)) > preprocess:
                    file_list = file_list[int(preprocess):]
                    self.preprocess = int(preprocess)
        # 创建窗口、画布，设置大小，设置字体等
        window = Tk()
        window.title('Plotting in Tkinter')
        window.geometry("1920x1080")
        # 设置图的字体大小
        plt.rcParams.update({'font.size': 30})
        fig = figure.Figure(figsize=(60, 20), dpi=35)
        plot1 = fig.add_subplot(111)
        plot1.plot()
        canvas = Canvas(window, width=1920, height=720, bg='white')
        output = FigureCanvasTkAgg(fig, master=canvas)
        output.draw()
        output.get_tk_widget().pack()
        canvas.pack()
        # 获取一个数据的生成器
        def get_one_data():
            for file_name in file_list:
               one_data =  pd.read_csv(f'日线_成交量_星期_data/raw_data/{file_name}')
               yield file_name,one_data
        next_data = get_one_data()
        self.last_filename,self.file_data = next(next_data)
        self.plt_one_data(self.file_data,train_days,plot1,output)
        # 按钮点击函数1
        def next_noe_1():
            self.save_noe_data(self.last_filename,self.file_data,1,self.preprocess)
            self.preprocess += 1
            self.last_filename,self.file_data = next(next_data)
            self.plt_one_data(self.file_data,train_days,plot1,output) # 绘制
        # 按钮点击函数1
        def next_noe_0():
            self.save_noe_data(self.last_filename,self.file_data,0,self.preprocess)
            self.preprocess += 1
            self.last_filename,self.file_data = next(next_data)
            self.plt_one_data(self.file_data,train_days,plot1,output) # 绘制
        label_1 = Button(master=window, command=next_noe_1, height=2, width=10, text="1")
        label_0 = Button(master=window, command=next_noe_0, height=2, width=10, text="0")
        # place the button
        label_1.pack()
        label_0.pack()
        # run the gui
        window.mainloop()
        
    def mkdata_by_auto(self,train_days=20,prd = False):
        # 检查有没有数据处理
        if not prd:
            raw_data = '日线_成交量_星期_data/raw_data'
        else:
            raw_data = '日线_成交量_星期_data/prd_data'
        if not path.exists(raw_data):
            print('没有数据可处理')
        file_list = os.listdir(raw_data)
        file_len = len(file_list)
        if file_len == 0:
            print('没有数据可处理')
        # 按照排序读取进度
        file_list.sort()
        self.preprocess = 0
        if not path.exists('日线_成交量_星期_data/train_data'):
            os.mkdir('日线_成交量_星期_data/train_data')
        else: # 获取之前的进度
            train_file = os.listdir('日线_成交量_星期_data/train_data')
            if not len(train_file) == 0:
                train_file.sort()
                preprocess = train_file[-1].split('.')[0]
                if preprocess.isdigit() and str(file_len) > preprocess:
                    file_list = file_list[int(preprocess):]
                    self.preprocess = int(preprocess)
        for file_name in file_list:
            one_data =  pd.read_csv(f'{raw_data}/{file_name}')
            if prd:
                target = file_name.split('_')[0]
            else:
                target = self.get_target(one_data,train_days)
            self.save_noe_data(file_name,one_data,target,self.preprocess)
            self.preprocess += 1
            print('\r正在处理：{}/{}'.format(self.preprocess,file_len),end='',flush=True)
        
    def download_data(self,data_days = 30,train_days = 20,data_space = 5,ref_days = 1):
        """
        配置详情
        1、data_days: 下载的训练数据的大小，默认30天
        2、train_days: 用于训练的数据大小，默认20天
        3、data_space: 取数据间隔，默认5天，则默认每只股票生成3个训练数据
        4、其他条件修改get_stock_basic参数
        5、ref_days : 参考涨跌结果的天数
        """
        # 获取主板上市公司股票代码
        code_list = self.data.get_stock_basic(market='主板',list_status='L')['ts_code']
        # 检查保存数据的目录
        if not path.exists('日线_成交量_星期_data'):
            os.mkdir('日线_成交量_星期_data')
        if not path.exists('日线_成交量_星期_data/raw_data'):
            os.mkdir('日线_成交量_星期_data/raw_data')
        # 显示进度
        fin = 0
        length = len(code_list)
        for code in code_list:
            # 下载数据
            data = self.data.get_by_code(code,data_days+ref_days).sort_values(by='trade_date',ascending=True)
            fin += 1
            if len(data) < train_days+ref_days:
                continue
            # 每隔data_space天保存一次
            for i in range(train_days,data_days+1,data_space):
                sub_data = data.iloc[i-train_days:i+ref_days]
                if len(sub_data) != train_days+ref_days:
                    break
                filename = '{}_{}.csv'.format(sub_data.iloc[-ref_days-1]['ts_code'],sub_data.iloc[-ref_days-1]['trade_date'])
                if not path.exists(f'日线_成交量_星期_data/raw_data/{filename}'):
                    sub_data.to_csv(f'日线_成交量_星期_data/raw_data/{filename}',index=False)
            print('\r正在下载 {} 数据, 进度: {}/{}'.format(code,fin,length),end='',flush=True)
    
    def download_predict_data_today(self,train_days = 20):
        # 获取主板上市公司股票代码
        code_list = self.data.get_stock_basic(market='主板',list_status='L')['ts_code']
        # 检查保存数据的目录
        if not path.exists('日线_成交量_星期_data'):
            os.mkdir('日线_成交量_星期_data')
        if not path.exists('日线_成交量_星期_data/prd_data'):
            os.mkdir('日线_成交量_星期_data/prd_data')
        # 显示进度
        fin = 0
        length = len(code_list)
        for code in code_list:
            # 下载数据
            data = self.data.get_by_code(code,train_days).sort_values(by='trade_date',ascending=True)
            fin += 1
            if len(data) < train_days:
                continue
            filename = '{}_{}.csv'.format(data.iloc[-1]['ts_code'],data.iloc[-1]['trade_date'])
            data.to_csv(f'日线_成交量_星期_data/prd_data/{filename}',index=False)
            print('\r正在下载 {} 数据, 进度: {}/{}'.format(code,fin,length),end='',flush=True)